
    /**  
    * @Title: ChooseDataSource.java
    * @Package com.wq.datasource3
    * @Description: 
    * @author likai
    * @date 2017年4月25日
    * @version V1.0  
    */
    
package com.frames.datasource;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

import javax.sql.DataSource;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;


/**
    * @ClassName: ChooseDataSource
    * @Description: 数据库动态切换
    * @author likai
    * @date 2017年4月25日
    *
    */

public class ChooseDataSource extends AbstractRoutingDataSource {
	 private static final ThreadLocal<String> holder = new ThreadLocal<String>();  //为了线程安全
	 private DataSource write; // 写库，只允许有一个
	 private List<DataSource> reads; // 读库，允许有多个
	 private AtomicLong readCount = new AtomicLong();//加锁机制的long
	 private int readSize = 0;//读库数量
	 private Map<Object, Object> dataSources = new HashMap<Object, Object>();//targetDataSources
	 public static final String WRITE = "WRITE";
	 public static final String READ = "READ";
	  
		 /**
		     * 初始化
		     */
	    @Override
	    public void afterPropertiesSet() {
	        if (write == null) {
	            throw new IllegalArgumentException("写库write必须配置");
	        }
	        if (reads.size()<1) {
	        	 throw new IllegalArgumentException("读库read至少一个");
			}
	        dataSources.put(WRITE, write);
	        if (null != reads && reads.size() > 0) {
	            for (int i = 0; i < reads.size(); i++) {
	                dataSources.put(READ + (i + 1), reads.get(i));
	            }
	            readSize = reads.size();
	        }
	        //this.setDefaultTargetDataSource(WRITE);
	        this.setTargetDataSources(dataSources);
	        super.afterPropertiesSet();
	    }
	
	 /**
     * 如果是选择使用从库，且从库的数量大于1，则通过取模来控制从库的负载,
     * 计算结果返回AbstractRoutingDataSource
     */
    @Override
    protected Object determineCurrentLookupKey() {   	
    	String key=getDataSource();//获取当前线程选择的数据库
    	if (WRITE.equals(key)) {
        	//写库
            return key;
        } else {
        	//读库根据通过取模来路由
            if (readSize > 1) {
                long c = readCount.incrementAndGet();
                c = c % readSize;
                return READ + (c + 1);
            } else {
                return READ + "1";
            }
        }
    }
	
	
	  /**  
     * 绑定当前线程数据源路由的key     
     * @param key  
     */  
	public static void putDataSource(String datasource) {  
		holder.set(datasource); 
	}  
      
    /**  
     * 获取当前线程的数据源路由的key  
     * @return  
     */  
	public static String getDataSource() {  
	    return holder.get();  
	}

	public DataSource getWrite() {
		return write;
	}

	public void setWrite(DataSource write) {
		this.write = write;
	}

	public List<DataSource> getReads() {
		return reads;
	}

	public void setReads(List<DataSource> reads) {
		this.reads = reads;
	}

	
	  
}
